from __future__ import print_function
from builtins import range
import os
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
from plyfile import PlyData, PlyElement, make2d, PlyParseError, PlyProperty
import numpy as np
import h5py

SAMPLING_BIN = os.path.join(BASE_DIR, "third_party/mesh_sampling/build/pcsample")

SAMPLING_POINT_NUM = 2048
SAMPLING_LEAF_SIZE = 0.005

MODELNET40_PATH = "../datasets/modelnet40"


def export_ply(pc, filename):
    vertex = np.zeros(pc.shape[0], dtype=[("x", "f4"), ("y", "f4"), ("z", "f4")])
    for i in range(pc.shape[0]):
        vertex[i] = (pc[i][0], pc[i][1], pc[i][2])
    ply_out = PlyData([PlyElement.describe(vertex, "vertex", comments=["vertices"])])
    ply_out.write(filename)


# Sample points on the obj shape
def get_sampling_command(obj_filename, ply_filename):
    cmd = SAMPLING_BIN + " " + obj_filename
    cmd += " " + ply_filename
    cmd += " -n_samples %d " % SAMPLING_POINT_NUM
    cmd += " -leaf_size %f " % SAMPLING_LEAF_SIZE
    return cmd


# --------------------------------------------------------------
# Following are the helper functions to load MODELNET40 shapes
# --------------------------------------------------------------

# Read in the list of categories in MODELNET40
def get_category_names():
    shape_names_file = os.path.join(MODELNET40_PATH, "shape_names.txt")
    shape_names = [line.rstrip() for line in open(shape_names_file)]
    return shape_names


# Return all the filepaths for the shapes in MODELNET40
def get_obj_filenames():
    obj_filelist_file = os.path.join(MODELNET40_PATH, "filelist.txt")
    obj_filenames = [
        os.path.join(MODELNET40_PATH, line.rstrip()) for line in open(obj_filelist_file)
    ]
    print("Got %d obj files in modelnet40." % len(obj_filenames))
    return obj_filenames


# Helper function to create the father folder and all subdir folders if not exist
def batch_mkdir(output_folder, subdir_list):
    if not os.path.exists(output_folder):
        os.mkdir(output_folder)
    for subdir in subdir_list:
        if not os.path.exists(os.path.join(output_folder, subdir)):
            os.mkdir(os.path.join(output_folder, subdir))


# ----------------------------------------------------------------
# Following are the helper functions to load save/load HDF5 files
# ----------------------------------------------------------------

# Write numpy array data and label to h5_filename
def save_h5_data_label_normal(
    h5_filename,
    data,
    label,
    normal,
    data_dtype="float32",
    label_dtype="uint8",
    normal_dtype="float32",
):
    h5_fout = h5py.File(h5_filename)
    h5_fout.create_dataset(
        "data", data=data, compression="gzip", compression_opts=4, dtype=data_dtype
    )
    h5_fout.create_dataset(
        "normal",
        data=normal,
        compression="gzip",
        compression_opts=4,
        dtype=normal_dtype,
    )
    h5_fout.create_dataset(
        "label", data=label, compression="gzip", compression_opts=1, dtype=label_dtype
    )
    h5_fout.close()


# Write numpy array data to h5_filename
def save_h5_data(h5_filename, data, data_dtype="float32"):
    h5_fout = h5py.File(h5_filename, "w")
    h5_fout.create_dataset(
        "data", data=data, compression="gzip", compression_opts=4, dtype=data_dtype
    )
    h5_fout.close()


# Write numpy array label to h5_filename
def save_h5_label(h5_filename, label, label_dtype="uint8"):
    h5_fout = h5py.File(h5_filename, "w")
    h5_fout.create_dataset(
        "label", data=label, compression="gzip", compression_opts=1, dtype=label_dtype
    )
    h5_fout.close()


# Write numpy array data and label to h5_filename
def save_h5(h5_filename, data, label, data_dtype="uint8", label_dtype="uint8"):
    h5_fout = h5py.File(h5_filename, "w")
    h5_fout.create_dataset(
        "data", data=data, compression="gzip", compression_opts=4, dtype=data_dtype
    )
    h5_fout.create_dataset(
        "label", data=label, compression="gzip", compression_opts=1, dtype=label_dtype
    )
    h5_fout.close()


# Read numpy array data and label from h5_filename
def load_h5_data_label_normal(h5_filename):
    f = h5py.File(h5_filename)
    data = f["data"][:]
    label = f["label"][:]
    normal = f["normal"][:]
    return (data, label, normal)


# Read numpy array data and label from h5_filename
def load_h5_data_label_seg(h5_filename):
    f = h5py.File(h5_filename)
    data = f["data"][:]
    label = f["label"][:]
    seg = f["pid"][:]
    return (data, label, seg)


# Read numpy array label from h5_filename
def load_h5_data(h5_filename):
    f = h5py.File(h5_filename)
    data = f["data"][:]
    return data


# Read numpy array label from h5_filename
def load_h5_label(h5_filename):
    f = h5py.File(h5_filename)
    label = f["label"][:]
    return label


# Read numpy array data and label from h5_filename
def load_h5(h5_filename):
    f = h5py.File(h5_filename)
    data = f["data"][:]
    label = f["label"][:]
    return (data, label)


# ----------------------------------------------------------------
# Following are the helper functions to load save/load PLY files
# ----------------------------------------------------------------

# Load PLY file
def load_ply_data(filename, point_num):
    plydata = PlyData.read(filename)
    pc = plydata["vertex"].data[:point_num]
    pc_array = np.array([[x, y, z] for x, y, z in pc])
    return pc_array


# Load PLY file
def load_ply_normal(filename, point_num):
    plydata = PlyData.read(filename)
    pc = plydata["normal"].data[:point_num]
    pc_array = np.array([[x, y, z] for x, y, z in pc])
    return pc_array


# Make up rows for Nxk array
# Input Pad is 'edge' or 'constant'
def pad_arr_rows(arr, row, pad="edge"):
    assert len(arr.shape) == 2
    assert arr.shape[0] <= row
    assert pad == "edge" or pad == "constant"
    if arr.shape[0] == row:
        return arr
    if pad == "edge":
        return np.lib.pad(arr, ((0, row - arr.shape[0]), (0, 0)), "edge")
    if pad == "constant":
        return np.lib.pad(arr, ((0, row - arr.shape[0]), (0, 0)), "constant", (0, 0))
